import numpy
import rknnlite
from rknnlite.api import RKNNLite
import cv2
import platform
import time

def array2bytes(array_img, suffix):
    # coding img
    success, encoded_array = cv2.imencode("." + suffix, array_img)
    # to bytes
    bytes_img = encoded_array.tobytes()

    return bytes_img

# decice tree for rk356x/rk3588
DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'

def get_host():
    # get platform and device type
    system = platform.system()
    machine = platform.machine()
    os_machine = system + '-' + machine
    if os_machine == 'Linux-aarch64':
        try:
            with open(DEVICE_COMPATIBLE_NODE) as f:
                device_compatible_str = f.read()
                if 'rk3588' in device_compatible_str:
                    host = 'RK3588'
                elif 'rk3562' in device_compatible_str:
                    host = 'RK3562'
                else:
                    host = 'RK3566_RK3568'
        except IOError:
            print('Read device node {} failed.'.format(DEVICE_COMPATIBLE_NODE))
            exit(-1)
    else:
        host = os_machine
    return host

RK3566_RK3568_RKNN_MODEL = 'resnet18_for_rk3566_rk3568.rknn'
RK3588_RKNN_MODEL = 'export_funds_stillgan.rknn'
RK3562_RKNN_MODEL = 'resnet18_for_rk3562.rknn'


if __name__ == '__main__':

    host_name = get_host()
    if host_name == 'RK3566_RK3568':
        rknn_model = RK3566_RK3568_RKNN_MODEL
    elif host_name == 'RK3562':
        rknn_model = RK3562_RKNN_MODEL
    elif host_name == 'RK3588':
        rknn_model = RK3588_RKNN_MODEL
    else:
        print("This demo cannot run on the current platform: {}".format(host_name))
        exit(-1)

    rknn_lite = RKNNLite(verbose=True)

    t1 = time.perf_counter()
    ret = rknn_lite.load_rknn(rknn_model)


    # run on RK356x/RK3588 with Debian OS, do not need specify target.
    if host_name == 'RK3588':
        ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
    else:
        ret = rknn_lite.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')
    t2 = time.perf_counter()


    ori_img = cv2.imread("test.jpg")
    # preprocess
    ori_size = ori_img.shape[:2]

    trans_image = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
    trans_image = cv2.resize(trans_image, (512, 512))
    trans_image = trans_image / 255.0
    img_arr = (trans_image - 0.5) / 0.5

    x = numpy.expand_dims(img_arr.astype(numpy.float32).transpose((2, 0, 1)), axis=0)

    outputs = rknn_lite.inference(inputs=[x, ])

    H, W = ori_size
    image = (numpy.transpose(outputs[0][0], (1, 2, 0)) + 1) / 2.0 * 255.0
    (r, g, b) = cv2.split(image)
    img_arr = cv2.merge([b, g, r])
    img_arr = cv2.resize(img_arr, (W, H), interpolation=cv2.INTER_CUBIC)
    img_bytes = array2bytes(img_arr, "png")
    image = numpy.asarray(bytearray(img_bytes), dtype=numpy.uint8)
    img_arr = cv2.imdecode(image, cv2.IMREAD_COLOR)

    cv2.imwrite("out.png", img_arr)
    t3 = time.perf_counter()
    print(f"init cost:{t2 - t1}\ninfer cost:{t3 - t2}")


    rknn_lite.release()